import os
import math
import time
import json
import numpy as np
import pandas as pd


def now():
    return int(time.time() * 1000)


def mapping_relationship(arr1, arr2):
    """
    Desc
    ----------
    为两个给定数组中的元素，建立映射关系。具体来说，将两个
    数组中的元素按照下标顺序依次排列在映射数组中，对于有歧
    义的映射对，使用传递关系消除歧义。

    Parameters
    ----------
    arr1:
        给定的整数数组1，长度与2相同
    arr2:
        给定的整数数组2，长度与1相同

    Returns
    ----------
    返回一个一维整数数组，顺序的每两个整数组成一对映射

    Example
    ----------
    >>> mapping_relationship([1, 2, 3, 4], [5, 6, 7, 8])
    [1, 5, 2, 6, 3, 7, 4, 8]

    >>> mapping_relationship([2, 3, 4, 5], [5, 8, 7, 1])
    [2, 1, 3, 8, 4, 7]

    >>> mapping_relationship([5, 8, 7, 1], [2, 3, 4, 5])
    [1, 2, 8, 3, 7, 4]
    """
    # 1. 数据验证
    size = len(arr1)
    assert size == len(arr2)
    mapping = np.full(size * 2, -1)
    # 2. 遍历数组，建立映射关系数组
    for i in range(size):
        e1 = arr1[i]
        e2 = arr2[i]
        idx1 = index_of(mapping, e1)
        idx2 = index_of(mapping, e2)
        # 2.1. 已有重复映射（关于 e1）
        if idx1 != -1:
            mapping[idx1] = e2
        # 2.2. 已有重复映射（关于 e2）
        elif idx2 != -1:
            mapping[idx2] = e1
        # 2.3. 不存在重复映射
        else:
            mapping[i * 2] = e1
            mapping[i * 2 + 1] = e2
        # 2.4. 添加完成后又出现重复（此时 e1 重复已被处理，即 e2 已被填入 mapping[idx1]，还需处理 e2 重复）
        #      e.g. [idx2, idx2+1, ..., idx1-2, idx1-1, idx1 -> e2] ---> [idx1-1, idx2+1, ..., idx1-2]
        if (idx1 != -1) and (idx2 != -1):
            if idx1 > idx2:
                mapping[idx2] = mapping[idx1 - 1]
                mapping[idx1 - 1], mapping[idx1] = -1, -1
            else:
                mapping[idx1] = mapping[idx2 + 1]
                mapping[idx2], mapping[idx2 + 1] = -1, -1
    # 3. 返回建立完成的映射数组
    return mapping


def index_of(arr, e):
    """
    查找给定数组中指定元素的下标，如果存在返回该下标（整数），
    如果不存在返回 -1

        Parameters
        ----------
        arr:
            给定的数组
        e:
            需要查找的元素

        Returns
        ----------
        如果存在返回指定元素在 np_arr 中的对应下标（整数），如果不存在返回 -1
    """
    res = np.where(np.array(arr) == e)[0]
    idx = -1
    if len(res) != 0:
        idx = res[0]
    return idx


def cardano(p, r):
    """
    计算以下形式的一元三次方程的根
        x^3 + px^2 + r = 0
    1. 等价变形为缺项形式（y = x + p/3）
        y^3 - (p^2/3)y + (p^3/9 - p^2/27 + r) = 0
    2. 根据 Cardano 公式得到方程的根
    """
    p_ = (- 1 / 3) * (p ** 2)
    r_ = ((p ** 3) / 9) - ((p ** 2) / 27) + r
    if (((p_ ** 3) / 27) + ((r_ ** 2) / 4)) < 0:
        return 0
    alpha = math.cbrt((-r / 2) + math.sqrt(((p_ ** 3) / 27) + ((r_ ** 2) / 4)))
    beta = math.cbrt((-r / 2) - math.sqrt(((p_ ** 3) / 27) + ((r_ ** 2) / 4)))
    y = alpha + beta
    x = y - (p / 3)
    return x


def sheng_jin_formula(a, b, c, d):
    delta_a = (b ** 2) - (3 * a * c)
    delta_b = (b * c) - (9 * a * d)
    delta_c = (c ** 2) - (3 * b * d)
    # 1. 方程有一个三重实根
    if delta_a == delta_b == delta_c == 0:
        return [-b / (3 * a)]
    delta = delta_b ** 2 - 4 * delta_a * delta_c
    # 2. 方程有一实根，一对虚根
    if delta > 0:
        y1 = delta_a * b + 3 * a * (-delta_b + math.sqrt((delta_b ** 2) - (4 * delta_a * delta_c))) / 2
        y2 = delta_a * b + 3 * a * (-delta_b - math.sqrt((delta_b ** 2) - (4 * delta_a * delta_c))) / 2
        x = (-b - (math.cbrt(y1) + math.cbrt(y2))) / (3 * a)
        return [x]
    # 2. 方程有三实根，其中包含一单根，一重根
    if delta == 0:
        k = delta_b / delta_a
        x1 = (-b / a) + k
        x2 = -k / 2
        return [x1, x2]
    # 3. 方程有三不等实根
    if delta < 0 < delta_a:
        t = (2 * delta_a * b - 3 * delta_b * a) / (2 * math.sqrt(delta_a ** 3))
        theta = np.arccos(t)
        x1 = (-b - 2 * math.sqrt(delta_a) * math.cos(theta / 3)) / (3 * a)
        x2 = (-b + math.sqrt(delta_a) * (math.cos(theta / 3) + math.sqrt(3) * math.sin(theta / 3))) / (3 * a)
        x3 = (-b + math.sqrt(delta_a) * (math.cos(theta / 3) - math.sqrt(3) * math.sin(theta / 3))) / (3 * a)
        return [x1, x2, x3]


class DataCollector:

    def __init__(self):
        self.datas = []
        self.data = []

    def collect(self, *args):
        if args:
            for arg in args:
                self.data.append(arg)

    def next(self):
        self.datas.append(self.data)
        self.data = []

    def write(self, file):
        df = pd.DataFrame(self.datas)
        file = os.path.join(os.path.dirname(os.getcwd()), f'data/experiment/{file}')
        df.to_csv(file, mode='a', header=False, index=False)

    @classmethod
    def collect_part(cls, env):
        parts = []
        for i in range(env.G):
            part = {'center': env.parts[i].pos.tolist(), 'ue': env._ues[env.parts_ues[i]].tolist()}
            parts.append(part)
        file = os.path.join(os.path.dirname(os.getcwd()), f'data/experiment/data_02_01.json')
        with open(file, 'w') as f:
            json.dump(parts, f)
